package com.gmrz.webauthn.handler;

import com.gmrz.service.challengestore.ChallengeStoreService;
import com.gmrz.uaf.common.Constants;
import com.gmrz.uaf.common.GuiceUtil;
import com.gmrz.uaf.crypto.AlgorithmNotFoundException;
import com.gmrz.uaf.crypto.CryptoEngine;
import com.gmrz.uaf.crypto.spi.UAFAuthAlgorithmSuite;
import com.gmrz.uaf.db.DAOException;
import com.gmrz.uaf.db.UAFDAOFactory;
import com.gmrz.uaf.db.UAFDBConnectionMgr;
import com.gmrz.uaf.protocol.v1.processor.exception.SignatureValidationException;
import com.gmrz.uaf.protocol.v1.processor.exception.UAFErrorCode;
import com.gmrz.uaf.servlet.v1.handler.UAFBaseHandler;
import com.gmrz.uaf.servlet.v1.validation.PayloadValidationException;
import com.gmrz.uaf.servlet.v1.validation.RestValidator;
import com.gmrz.util.Convert;
import com.gmrz.util.CryUtil;
import com.gmrz.util.Streams;
import com.gmrz.util.Strings;
import com.gmrz.util.db.DBUtil;
import com.gmrz.webauthn.common.WebAuthnConstants;
import com.gmrz.webauthn.db.dao.WebAuthnAuthenticatorsDAO;
import com.gmrz.webauthn.protocol.v1.json.WebAuthnSchemaBuilder;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnErrorCode;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnException;
import com.gmrz.webauthn.protocol.v1.schema.*;
import com.gmrz.webauthn.request.*;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import org.apache.commons.codec.binary.Base64;
import org.apache.logging.log4j.Logger;
import org.apache.logging.log4j.LogManager;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.lang.reflect.Type;
import java.nio.charset.Charset;
import java.security.*;
import java.sql.Connection;
import java.util.Arrays;

public class WebAuthResponseHandler extends UAFBaseHandler {

    private static final Logger LOG = LogManager.getLogger(WebAuthResponseHandler.class);
    private CryptoEngine cryptoEngine;

    public WebAuthResponseHandler() {
        this.cryptoEngine = GuiceUtil.getProcessorInjector().getInstance(CryptoEngine.class);
    }

    public String authFinish(String request) {
        LOG.info("[/auth/send][begin]");
        Gson gson = WebAuthnSchemaBuilder.getGson();
        WebAuthnRespPayload<WebAuthnRequestContext> respPayload = new WebAuthnRespPayload<WebAuthnRequestContext>();
        Type reqType = new TypeToken<WebAuthnReqPayload<WebAuthnAuthContext>>() {
        }.getType();

        String response;
        Connection conn = null;
        try {
            WebAuthnReqPayload<WebAuthnAuthContext> reqPayload = gson.fromJson(request, reqType);
            //日志记录
            logRequest(request, reqPayload.getContext().getTransNo(), reqPayload.getContext().getUserName(),
                    reqPayload.getContext().getAppID(), reqPayload.getContext().getAuthType().get(0),reqPayload.getContext().getTransType(), Constants.Method.AUTH_FINISH);

            //验证请求参数
            validate(reqPayload);

            String rpID = WebAuthnConstants.getTenantConfigContent(WebAuthnConstants.WEBAUTHN_AUTHENTICATOR_RPID, reqPayload.getContext().getAppID());
            String rpName = WebAuthnConstants.getTenantConfigContent(WebAuthnConstants.WEBAUTHN_AUTHENTICATOR_RPNAME, reqPayload.getContext().getAppID());

            String rpIDSha256 = Convert.toBase64(CryUtil.getSHA256(rpID.getBytes()));


            String clientDataStr = reqPayload.getCredentials().getResponse().getClientDataJSON();
            String clientDataDecodedStr = new String(Base64.decodeBase64(clientDataStr), Charset.forName(Constants.UTF_8));
            CollectedClientData clientData = WebAuthnSchemaBuilder.getGson().fromJson(clientDataDecodedStr, CollectedClientData.class);

            String serverData = new String(Convert.fromBase64(reqPayload.getServerData()));
            Type optionsType = new TypeToken<PublicKeyCredentialRequestOptions>() {
            }.getType();
            PublicKeyCredentialRequestOptions options = gson.fromJson(serverData, optionsType);

            //对报文中的挑战值进行校验
            ChallengeStoreService challengeStoreService = GuiceUtil.getProcessorInjector().getInstance(ChallengeStoreService.class);
            String challengeContent = challengeStoreService.getRedisValue(options.getChallenge());
            if (Strings.isNullOrEmpty(challengeContent)) {
                throw new WebAuthnException(WebAuthnErrorCode.STATUS_CHALLENGE_VALIDATION_FAILED);
            }

            AuthenticatorAssertionResponse authenticatorResponse = GuiceUtil.getProcessorInjector().getInstance(AuthenticatorAssertionResponse.class);
            //签名数据
            authenticatorResponse.setSignature(Convert.strictBase64URLdecode(reqPayload.getCredentials().getResponse().getSignature()));

            //客户端数据
            authenticatorResponse.setClientDataBytes(Convert.strictBase64URLdecode(clientDataStr));
            authenticatorResponse.setClientData(clientData);

            byte[] authenticatorDataBytes = Convert.strictBase64URLdecode(reqPayload.getCredentials().getResponse().getAuthenticatorData());
            authenticatorResponse.setAuthenticatorDataBytes(authenticatorDataBytes);

            AuthenticatorData authData = new AuthenticatorData();
            authData.decode(authenticatorResponse.getAuthenticatorDataBytes());

            authenticatorResponse.setAuthenticatorData(authData);

            //UserHandle
            String userHandleStr = reqPayload.getCredentials().getResponse().getUserHandle();
            byte[] userHandle = null;
            if (Strings.isNotBlank(userHandleStr)) {
                userHandle = userHandleStr.getBytes(Charset.forName("UTF-8"));
            }
            authenticatorResponse.setUserHandle(userHandle);

            conn = UAFDBConnectionMgr.getConnection(true);
            WebAuthnAuthenticatorsDAO webauthnAuthenticatorsDAO = UAFDAOFactory.createWebAuthnAuthenticatorsDAO(conn);
            WebAuthenticator authenticator = webauthnAuthenticatorsDAO.findByRawIdAndTenantId(reqPayload.getCredentials().getId(), reqPayload.getContext().getAppID());


            if (authenticator == null) {
                throw new DAOException(UAFErrorCode.DB_QUERY_AUTHENTICATOR_FAILED);
            }

            AuthenticatorData authenticatorData = authenticatorResponse.getAuthenticatorData();
            String authDataRpIDSha256 = Convert.toBase64(authenticatorData.getRpID());
            if(!rpIDSha256.equals(authDataRpIDSha256)){
                LOG.error("Failed to validate rpID hash in AuthenticatorData. AuthDataRpIDSha256:"+authDataRpIDSha256+",rpIDSha256:"+rpIDSha256);
                throw new WebAuthnException(WebAuthnErrorCode.STATUS_RPID_VALIDATION_FAILED);
            }
            //验证认证次数
            if (authenticator.getSignCounter() >= authenticatorResponse.getAuthenticatorData().getSignCount()) {
                throw new SignatureValidationException(UAFErrorCode.PROTOCOL_INVALID_SIGN_COUNTER);
            }

            if (!validateUserHandle(authenticator, authenticatorResponse.getUserHandle())) {
                throw new SignatureValidationException(UAFErrorCode.PROTOCOL_SIGNDATA_VALIDATION_FAILED);
            }


            byte[] signedData = getSignedData(authenticatorResponse);
            byte[] pubKey = Convert.fromBase64(authenticator.getPublicKey());

            //验签
            boolean verifySignatureFlag = validateSignature(authenticator.getKeyFormat(), pubKey, authenticator.getSignAlgo(), signedData, authenticatorResponse.getSignature());

            if (!verifySignatureFlag) {
                throw new SignatureValidationException(UAFErrorCode.PROTOCOL_SIGNDATA_VALIDATION_FAILED);
            }
            String hashedUserName = generateHashedUserName(reqPayload.getContext().getUserName(), reqPayload.getContext().getRf1(), reqPayload.getContext().getRf2());

            boolean updateDbFlag = webauthnAuthenticatorsDAO.updateAuthCounter(hashedUserName, reqPayload.getCredentials().getId(), reqPayload.getContext().getAppID(), authenticatorResponse.getAuthenticatorData().getSignCount());

            if (!updateDbFlag) {
                throw new DAOException(UAFErrorCode.PROTOCOL_SIGNDATA_VALIDATION_FAILED);
            }

            respPayload.setStatusCode(Constants.StatusCode.OK.getCode());
            respPayload.setDescription("OK");
            response = gson.toJson(respPayload);

        } catch (Exception ex) {
            LOG.error("[/reg/send] ERROR", ex);
            response = buildWebAuthnRespFromException(respPayload, ex);
        }finally {
            DBUtil.close(conn);
        }

        LOG.info("[/auth/send][end]");
        return response;
    }


    private byte[] getSignedData(AuthenticatorAssertionResponse response)
            throws IOException, NoSuchAlgorithmException, WebAuthnException {
        ByteArrayOutputStream outputStream = null;
        try {
            outputStream = new ByteArrayOutputStream();


            outputStream.write(response.getAuthenticatorDataBytes());


            byte[] clientDataHash = getClientDataHash(response.getClientDataBytes(), response.getClientData().getHashAlg());

            LOG.debug(" client data hash size : {}", clientDataHash.length);
            outputStream.write(clientDataHash);

            return outputStream.toByteArray();
        } finally {
            Streams.close(outputStream);
        }
    }


    public byte[] getClientDataHash(byte[] clientDataBytes, HashAlgorithmIdentifier algorithm) {
        byte[] clientDataHash = null;
        switch (algorithm) {
            case S256:
                clientDataHash = CryUtil.getSHA256(clientDataBytes);
                break;
            default:
                break;
        }

        return clientDataHash;
    }


    public boolean validateUserHandle(WebAuthenticator authenticator, byte[] userHandle) {
        boolean flag = true;
        if ((userHandle != null) && (userHandle.length > 0)) {
            byte[] userHandleInAuthenticator = authenticator.getUserId().getBytes(Charset.forName("UTF-8"));

            if (!Arrays.equals(userHandleInAuthenticator, userHandle)) {
                flag = false;
            }
            LOG.debug(" completed userHandle validation");
        }
        return flag;
    }


    public boolean validateSignature(int keyFormat, byte[] pubKey, int signAlgo, byte[] signedData, byte[] signature) throws GeneralSecurityException, AlgorithmNotFoundException {
        PublicKey key = this.cryptoEngine.getKeyFormat(keyFormat).getKey(pubKey, signAlgo);
        UAFAuthAlgorithmSuite authAlgorithmSuite = this.cryptoEngine.getAuthAlgorithmSuite(signAlgo);
        return authAlgorithmSuite.verifySignature(key, signedData, signature);
    }

    private void validate(WebAuthnReqPayload<WebAuthnAuthContext> reqPayload) throws PayloadValidationException, UnsupportedEncodingException {
        if ((reqPayload == null) || (reqPayload.getContext() == null)) {
            throw new PayloadValidationException("context", RestValidator.VALIDATION_ERROR.EMPTY);
        }
        RestValidator.validate("userName", reqPayload.getContext().getUserName());
        RestValidator.validate("appID", reqPayload.getContext().getAppID());
        //RestValidator.validate("deviceID", reqPayload.getContext().getDevices().getDeviceID());
        RestValidator.validate("authType", reqPayload.getContext().getAuthType());
        RestValidator.validate("transType", reqPayload.getContext().getTransType());
        RestValidator.validate("transNo", reqPayload.getContext().getTransNo());
        RestValidator.validate("rf1", reqPayload.getContext().getRf1());
        RestValidator.validate("rf2", reqPayload.getContext().getRf2());
        RestValidator.validate("protocol", reqPayload.getContext().getProtocol());
    }
}
